from fastapi import APIRouter, WebSocket, WebSocketDisconnect
import json
import asyncio
from datetime import datetime
from process import get_easyquotation_sina
from controller import CrealMarketStatic
from process import baseFunction as bs

router = APIRouter()
clients = []


def get_current_time():
    return datetime.now()


@router.websocket("/ws1")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            now = get_current_time()
            data = 'hello'

            print(f"Received message: {data} {now}")
            await websocket.send_text(f"Message received: {data} {now}")
            await asyncio.sleep(15)

    except WebSocketDisconnect:
        print("WebSocket disconnected")
    finally:
        await websocket.close()


@router.websocket("/zhishu")
async def websocket_stock_data(websocket: WebSocket):
    await websocket.accept()
    try:
        async for stock_data in get_easyquotation_sina.fetch_stock_data():
            await websocket.send_text(stock_data)
    except Exception as e:
        error_message = {"指数行情获取失败": str(e)}
        await websocket.send_text(json.dumps(error_message))
    finally:
        await websocket.close()


@router.websocket("/realMarket")
async def websocket_realMarket(websocket: WebSocket):
    await websocket.accept()
    clients.append(websocket)
    try:
        while True:
            # 在交易时间内发送市场数据
            if  bs.is_market_open():
                market_data = CrealMarketStatic.getRealMarketStatic()
                await websocket.send_text(market_data)

                await asyncio.sleep(5)  # 每5秒检查一次
            else:
                market_data = '非交易时间，没有及时行情数据。'
                await websocket.send_text(market_data)
                clients.remove(websocket)
                break

    except WebSocketDisconnect:
        clients.remove(websocket)
